import datetime
from typing import List, Optional, Tuple, Dict

from common.logger import logger
from common.typing_utils import numeric

from common.geometry.geom import Geom, BoundingBox
from common.geoimage.raster_dataset import RasterDataset
from common.geoimage.gdal_sys import GDALRaster, SourceWindow

from catalog.models.known_models import KNOWN_MODELS
from catalog.models.base_stac_item import BaseSTACItem

from catalog.raster.strategy_recent import composite_recent
from catalog.raster.strategy_recent_s2 import composite_recent_s2_v2
from catalog.raster.utils import get_s2_l2a_add_boa_offset, update_s2_l2a_boa
from catalog.raster.strategy_s2decloud_v2 import composite_s2decloud_v2


KNOWN_STRATEGY = {
    "recent": composite_recent,
    "recent_s2_v2": composite_recent_s2_v2,
    "s2decloud-v2": composite_s2decloud_v2,
}


def _rds_to_file(rds, fpath_dest):
    if fpath_dest.endswith(".jp2"):
        rds.to_jp2(fpath_dest)
    elif fpath_dest.endswith(".tif") or fpath_dest.endswith(".tiff"):
        rds.to_geotiff(fpath_dest)
    else:
        logger.warning(f"Unknown file type for {fpath_dest}, default to geotiff")
        rds.to_geotiff(fpath_dest)


def _check_validity(window, extent, aoi, dest_srs):
    n_provided = sum(item is not None for item in (window, extent, aoi))
    if n_provided > 1:
        raise ValueError("Should only provide one of `window`, `extent` or `aoi`.")

    if (extent is not None) and dest_srs is None:
        raise ValueError(
            "If `extent` is provided then `dest_srs` is required to be provided."
        )


def retrieve_by_item(
    item: BaseSTACItem,
    asset_names: List[str] = None,
    window: Optional[SourceWindow] = None,
    extent: Optional[BoundingBox] = None,
    aoi: Optional[Geom] = None,
    resample_alg: str = "near",
    crop_to_aoi: bool = False,
    dest_srs: Optional[str] = None,
    dest_res: Optional[numeric] = None,
    fpath_dest: Optional[str] = None,
) -> Optional[RasterDataset]:
    """Retrieve data of interest from a `BaseSTACItem` in our catalog.

    Parameters
    ----------
    item : BaseSTACItem
        A `BaseSTACItem` representing the data of interest in our catalog.
    asset_names : List[str], optional
        A list of asset names to be retrieved, by default None to use a default one for the corresponding collection.
        NOTE: We have not yet required each data collection to provide the default `asset_names`.
        So currently please provide the `asset_names`.
    window : Optional[SourceWindow], optional
        Output extent specified by source extent in image space coordinate, by default None.
    extent : Optional[BoundingBox], optional
        Output extent specified by coordinates under destination SRS, by default None
    aoi : Optional[Geom], optional
        Output extent specified by a geometry object, by default None
    resample_alg : str, optional
        Algorithm for resampling, if resampling is needed, by default "near".
        See <https://gdal.org/programs/gdalwarp.html#cmdoption-gdalwarp-r> for more algorithms supported.
    crop_to_aoi : bool, optional
        Whether to crop the retrieved data by the given AOI, by default False not to crop,
        notice it is only valid when the `aoi` argument is provided.
    dest_srs : Optional[str], optional
        Destination spatial reference system represented by an EPSG string, e.g. "EPSG:4326", "EPSG:4857", etc.
        by default None to use the item's own spatial reference system.
    dest_res : Optional[numeric], optional
        Destination spatial resolution, by default None to use the item's own spatial resolution.
    fpath_dest : Optional[str], optional
        File path to the retrieved raster,
        by default None so not to save a raster file but return the raster as a `RasterDataset` object.

    Returns
    -------
    Optional[RasterDataset]
        A `RasterDataset` object for the data retrieved,
        or `None` the specified item could not be found in our catalog,
        or `None` when a result raster file of geotiff or jpeg2000 format is requested to saved.

    Raises
    ------
    NotImplementedError
        Raised when `asset_names` is None.
    """
    _check_validity(window, extent, aoi, dest_srs)
    if asset_names is None:
        # TODO, grab a default asset_names; and each model should define its own default_asset_names
        raise NotImplementedError(
            "later we will have a default asset_names set in each model"
        )

    try:
        nodata_values = []
        for asset_name in asset_names:
            nodata_values.append(item.bands[asset_name].nodata_val)
        if len(nodata_values) == 1:
            nodata_values = nodata_values[0]
    except Exception:
        nodata_values = 0

    vrt_scene = item.choose_assets(asset_names, dest_nodata=nodata_values)
    if vrt_scene is None:
        logger.warning(f"Choose Assest Fail: {item.uid}. Bands: {asset_names}")
        return None

    vrt_scene_warped = vrt_scene.warp(
        dest_srs,
        dest_x_res=dest_res,
        dest_y_res=dest_res,
        resample_alg=resample_alg,
        dest_nodata=nodata_values,
    )
    with GDALRaster.open(vrt_scene_warped._vrt_file_path) as gdal_rst:
        rds = RasterDataset.from_gdal_raster(
            gdal_rst, window=window, extent=extent, aoi=aoi
        )
    if aoi is not None and crop_to_aoi:
        rds = rds.crop(aoi, nodata_values=nodata_values)

    # NOTE Post-process Sentinel-2 L2A due to ESA updates of -1000 offset
    asset_offsets = get_s2_l2a_add_boa_offset(item, asset_names)
    if asset_offsets is not None:
        rds.data = update_s2_l2a_boa(rds.data, asset_offsets)

    if fpath_dest is not None:
        _rds_to_file(rds, fpath_dest)
        return None
    else:
        return rds


def retrieve_by_uid(
    collection_id: str, uid: str, create_db_conn: bool = True, **kwargs
) -> Optional[RasterDataset]:
    """Retrieve data of interest from a `BaseSTACItem` in our catalog,
        with the item represented by a `collection_id` for the data collection it belongs to,
        and `uid` for itself in that data collection.

    Parameters
    ----------
    collection_id : str
        Data collection id for the data of interest, e.g "COPERNICUS/SENTINEL2/L2A"
    uid : str
        Unique id for the item of interest.
    **kwargs
        Additional argument passed to the `retrieve_by_item` function. Please see that function for more information.

    Returns
    -------
    Optional[RasterDataset]
        A `RasterDataset` object for the data retrieved,
        or `None` the specified item could not be found in our catalog,
        or `None` when a result raster file of geotiff or jpeg2000 format is requested to saved.
    """
    model = KNOWN_MODELS[collection_id]

    item = model.get_one_item(uid)

    if item is None:
        logger.warning(f"Cannot find the requested scene {uid} under {collection_id}")
        return None
    else:
        return retrieve_by_item(item, **kwargs)


def retrieve_by_items(
    items: List[BaseSTACItem],
    asset_names: List[str] = None,
    extent: Optional[BoundingBox] = None,
    aoi: Optional[Geom] = None,
    resample_alg: str = "near",
    crop_to_aoi: bool = False,
    dest_srs: Optional[str] = None,
    dest_res: Optional[numeric] = None,
    fpath_dest: Optional[str] = None,
    composite_strategy: Optional[str] = "recent",
    **kwargs,
) -> Optional[RasterDataset]:
    """Retrieve data of interest from a list of `BaseSTACItem`s in our catalog.

    Parameters
    ----------
    items : List[BaseSTACItem]
        A list of `BaseSTACItem`s representing the data of interest in our catalog.
    asset_names : List[str], optional
        A list of asset names to be retrieved, by default None to use a default one for the corresponding collection.
        NOTE: We have not yet required each data collection to provide the default `asset_names`.
        So currently please provide the `asset_names`.
    extent : Optional[BoundingBox], optional
        Output extent specified by coordinates under destination SRS, by default None
    aoi : Optional[Geom], optional
        Output extent specified by a geometry object, by default None
    resample_alg : str, optional
        Algorithm for resampling, if resampling is needed, by default "near".
        See <https://gdal.org/programs/gdalwarp.html#cmdoption-gdalwarp-r> for more algorithms supported.
    crop_to_aoi : bool, optional
        Whether to crop the retrieved data by the given AOI, by default False not to crop,
        notice it is only valid when the `aoi` argument is provided.
    dest_srs : Optional[str], optional
        Destination spatial reference system represented by an EPSG string, e.g. "EPSG:4326", "EPSG:4857", etc.
        by default None to use the item's own spatial reference system.
    dest_res : Optional[numeric], optional
        Destination spatial resolution, by default None to use the first item's own spatial resolution.
    fpath_dest : Optional[str], optional
        File path to the retrieved raster,
        by default None so not to save a raster file but return the raster as a `RasterDataset` object.
    composite_strategy : Optional[str], optional
        Composition strategy, could be None, or one of "recent" or "s2decloud-v1", by default None,
        the procedure would choose the strategy for you.
        For now if the items invovled are from Sentinel-2 L2A, "s2decloud-v1" would be used,
        otherwise, "recent" would be used.
    **kwargs
        Additional argument passed to the composition strategy specified by `composite_strategy`.
        Please see the corresponding function for more information.

    Returns
    -------
    Optional[RasterDataset]
        A `RasterDataset` object for the data retrieved,
        or `None` the specified item could not be found in our catalog,
        or `None` when a result raster file of geotiff or jpeg2000 format is requested to saved.

    Raises
    ------
    NotImplementedError
        Raised when `asset_names` is None.
    ValueError
        Raised when an unknown `composite_strategy` is provided.
    """
    _check_validity(None, extent, aoi, dest_srs)

    if asset_names is None:
        # TODO, grab a default asset_names; and each model should define its own default_asset_names
        raise NotImplementedError(
            "later we will have a default asset_names set in each model"
        )

    if (
        items[0].info["collection"] == "sentinel2-l2a"
        and composite_strategy == "recent"
    ):
        composite_strategy = "recent_s2_v2"

    if composite_strategy not in KNOWN_STRATEGY:
        raise ValueError(f"Unknown composition strategy {composite_strategy}.")

    composite_method = KNOWN_STRATEGY[composite_strategy]
    rds = composite_method(
        items,
        asset_names,
        extent,
        aoi,
        dest_srs,
        dest_res,
        resample_alg,
        **kwargs,
    )

    # TODO, we should decide the nodata_values based on stored information in the scene class
    # but for now, not all Scene class has that info, we should revisit and improve this part later
    try:
        nodata_values = []
        for asset_name in asset_names:
            nodata_values.append(items[0].bands[asset_name].nodata_val)
        if len(nodata_values) == 1:
            nodata_values = nodata_values[0]
    except Exception:
        nodata_values = 0

    if aoi is not None and crop_to_aoi:
        rds = rds.crop(aoi, nodata_values=nodata_values)
    if fpath_dest is not None:
        _rds_to_file(rds, fpath_dest)
        return None
    else:
        return rds


def retrieve_by_uids(
    collection_id: str,
    uids: List[str],
    composite_strategy: str = "recent",
    create_db_conn: bool = True,
    **kwargs,
) -> Optional[RasterDataset]:
    """Retrieve data of interest from a collection id and a list of unique item ids.

    Parameters
    ----------
    collection_id : str
        Data collection id for the data of interest, e.g "COPERNICUS/SENTINEL2/L2A"
    uids : List[str]
        A list of unique ids for the items of interest.
    **kwargs
        Additional argument passed to the `retrieve_by_items` function,
        or to the composition strategy specified by `composite_strategy`.
        Please see that function for more information.

    Returns
    -------
    Optional[RasterDataset]
        A `RasterDataset` object for the data retrieved,
        or `None` the specified item could not be found in our catalog,
        or `None` when a result raster file of geotiff or jpeg2000 format is requested to saved.
    """

    model = KNOWN_MODELS[collection_id]

    items = list(model.get_many_items(uids))
    items = [item for item in items if item is not None]

    n_items = len(items)
    if n_items < len(uids):
        logger.warning(f"Only {n_items} out of {len(uids)} could be found.")

    if n_items == 0:
        logger.warning("Cannot find any scene given.")
        return None
    else:
        return retrieve_by_items(items, composite_strategy=composite_strategy, **kwargs)


def retrieve_by_query(
    collection_id: str,
    spatial: Optional[Geom] = None,
    temporal: Optional[
        Tuple[Optional[datetime.datetime], Optional[datetime.datetime]]
    ] = None,
    properties: Optional[Dict] = None,
    asset_names: List[str] = None,
    resample_alg: str = "near",
    crop_to_aoi: bool = False,
    dest_srs: Optional[str] = None,
    dest_res: Optional[numeric] = None,
    fpath_dest: Optional[str] = None,
    composite_strategy: Optional[str] = "recent",
    create_db_conn: bool = True,
    **kwargs,
) -> Optional[RasterDataset]:
    """Retrieve data of interest from a collection id and query terms.

    Parameters
    ----------
    collection_id : str
        Data collection id for the data of interest, e.g "COPERNICUS/SENTINEL2/L2A"
    spatial : Optional[Geom], optional
        Spatial query as a `Geom` object that intersects with items of interest, by default None not to do spatial query
    temporal : Optional[Tuple[Optional[datetime.datetime], Optional[datetime.datetime]] ], optional
        Temporal query, as a tuple of start_datetime and end_datetime, by default None not to do temporal query
    properties : Optional[Dict], optional
        Property queries as dictionary, see each collection's model class for more details,
        by default None not to perform any property based query
    asset_names : List[str], optional
        A list of asset names to be retrieved, by default None to use a default one for the corresponding collection.
        NOTE: We have not yet required each data collection to provide the default `asset_names`.
        So currently please provide the `asset_names`.
    resample_alg : str, optional
        Algorithm for resampling, if resampling is needed, by default "near".
        See <https://gdal.org/programs/gdalwarp.html#cmdoption-gdalwarp-r> for more algorithms supported.
    crop_to_aoi : bool, optional
        Whether to crop the retrieved data by the given spatial, by default False not to crop.
    dest_srs : Optional[str], optional
        Destination spatial reference system represented by an EPSG string, e.g. "EPSG:4326", "EPSG:4857", etc.
        by default None to use the item's own spatial reference system.
    dest_res : Optional[numeric], optional
        Destination spatial resolution, by default None to use the first item's own spatial resolution.
    fpath_dest : Optional[str], optional
        File path to the retrieved raster,
        by default None so not to save a raster file but return the raster as a `RasterDataset` object.
    composite_strategy : Optional[str], optional
        Composition strategy, could be None, or one of "recent" or "s2decloud-v1", by default None,
        the procedure would choose the strategy for you.
        For now if the items invovled are from Sentinel-2 L2A, "s2decloud-v1" would be used,
        otherwise, "recent" would be used.
    create_db_conn : bool, optional
        Whether to create a new database connection in query, by default True to create a new one.
        Notice it is an internal argument reserved for internal developers to re-use existing database connection.

    Returns
    -------
    Optional[RasterDataset]
        A `RasterDataset` object for the data retrieved,
        or `None` the specified item could not be found in our catalog,
        or `None` when a result raster file of geotiff or jpeg2000 format is requested to saved.
    """
    model = KNOWN_MODELS[collection_id]

    model.init_connection()
    if properties is None:
        items = model.query_many_items(spatial=spatial, temporal=temporal)
    else:
        items = model.query_many_items(spatial=spatial, temporal=temporal, **properties)
    items = [item for item in items if item is not None]
    model.close_connection()

    n_items = len(items)
    logger.info(f"{n_items} items to be retrieved.")

    if n_items == 0:
        logger.warning("Cannot find any scene matching given query.")
        return None
    else:
        return retrieve_by_items(
            items,
            asset_names=asset_names,
            composite_strategy=composite_strategy,
            aoi=spatial,
            resample_alg=resample_alg,
            crop_to_aoi=crop_to_aoi,
            dest_srs=dest_srs,
            dest_res=dest_res,
            fpath_dest=fpath_dest,
        )
